#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI 安装脚本 (插件版本)
实现与 install_comfyui.bat 相同的功能
"""

import os
import sys
import subprocess
import shutil
import platform
import json
from pathlib import Path
from urllib.parse import urlparse
import requests
import time

# 导入环境配置 - 使用新的全局变量
from .env_setup import (
    PROJECT_ROOT, COMFYUI_DIR, PYTHON_DIR, PYTHON_CMD,
    get_python_cmd, get_project_root, get_comfyui_dir
)

class ComfyUIInstaller:
    """ComfyUI安装器插件类"""
    
    # 插件排序
    order = 1
    
    def __init__(self):
        """初始化安装器"""
        # 使用env_setup.py的全局变量
        self.workspace_path = PROJECT_ROOT  # 项目根目录
        self.comfyui_path = COMFYUI_DIR     # ComfyUI目录
        self.config_path = self.workspace_path / "Config"
        self.manager_path = self.config_path / "custom_nodes" / "ComfyUI-Manager"
        self.python_dir = PYTHON_DIR        # Python目录
        self.python_exe = self.python_dir / "python.exe"
        
        # 颜色定义
        self.RED = "\033[31m"
        self.GREEN = "\033[32m"
        self.YELLOW = "\033[33m"
        self.BLUE = "\033[34m"
        self.CYAN = "\033[36m"
        self.PURPLE = "\033[35m"
        self.RESET = "\033[0m"
    
    def start(self) -> bool:
        """启动安装器"""
        print(f"\n{self.PURPLE}■■■ ComfyUI 安装器 v1.0（插件版） ■■■{self.RESET}")
        print(f"工作目录: {self.CYAN}{self.workspace_path.absolute()}{self.RESET}")
        
        # 检查管理员权限
        if not self.check_admin():
            print(f"{self.RED}[错误] ComfyUI安装需要管理员权限!{self.RESET}")
            print("请右键'以管理员身份运行'")
            input("按回车键退出...")
            return False
        
        # 执行安装流程
        return self.install_comfyui()
    
    @staticmethod
    def Launch(self):
        """静态启动方法"""
        installer = ComfyUIInstaller()
        return installer.start()
    
    def run_command(self, cmd, shell=True, check=True, capture_output=True, text=True, show_output=True):
        """运行命令并返回结果"""
        try:
            if show_output:
                # 实时显示输出
                print(f"{self.CYAN}执行命令: {cmd}{self.RESET}")
                process = subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, 
                                         stderr=subprocess.STDOUT, text=text, bufsize=1, universal_newlines=True)
                
                # 实时读取输出
                output_lines = []
                while True:
                    if process.stdout is None:
                        break
                    output = process.stdout.readline()
                    if output == '' and process.poll() is not None:
                        break
                    if output:
                        output = output.strip()
                        if output:
                            print(f"  {output}")
                            output_lines.append(output)
                
                # 等待进程完成
                return_code = process.poll()
                if return_code is None:
                    return_code = -1
                result = subprocess.CompletedProcess(cmd, return_code, '\n'.join(output_lines), '')
                return result
            else:
                # 静默模式
                result = subprocess.run(cmd, shell=shell, check=check, 
                                      capture_output=capture_output, text=text)
                return result
        except subprocess.CalledProcessError as e:
            if check:
                print(f"{self.RED}[错误] 命令执行失败: {cmd}{self.RESET}")
                print(f"错误信息: {e.stderr}")
                return None
            else:
                # 当check=False时，返回异常对象作为结果
                return e
        except Exception as e:
            print(f"{self.RED}[错误] 命令执行异常: {cmd}{self.RESET}")
            print(f"异常信息: {e}")
            return None

    def check_admin(self):
        """检查管理员权限"""
        if platform.system() == "Windows":
            try:
                result = self.run_command("net session", check=False)
                return result is not None and result.returncode == 0
            except:
                return False
        return True

    def clone_repository(self, url, target_path, name):
        """克隆Git仓库"""
        if not target_path.exists():
            print(f"克隆 {name} 到 {target_path}")
            result = self.run_command(f'git clone --jobs=8 "{url}" "{target_path}"', check=False)
            if result and result.returncode == 0:
                print(f"{self.GREEN}✅ {name} 克隆成功{self.RESET}")
            else:
                print(f"{self.RED}❌ {name} 克隆失败{self.RESET}")
                return False
        else:
            print(f"{self.GREEN}✅ {name} 已存在{self.RESET}")
        return True

    def test_network_connection(self):
        """测试网络连接"""
        print(f"{self.BLUE}[网络连接测试]{self.RESET}")
        print("-" * 45)
        
        test_urls = [
            "https://pypi.org/simple/",
            "https://mirrors.cloud.tencent.com/pypi/simple/",
            "https://mirrors.aliyun.com/pypi/simple/",
            "https://pypi.tuna.tsinghua.edu.cn/simple/"
        ]
        
        working_urls = []
        for url in test_urls:
            try:
                response = requests.get(url, timeout=5)
                if response.status_code == 200:
                    print(f"{self.GREEN}✅ {url} - 连接正常{self.RESET}")
                    working_urls.append(url)
                else:
                    print(f"{self.YELLOW}⚠️ {url} - HTTP {response.status_code}{self.RESET}")
            except Exception as e:
                print(f"{self.RED}❌ {url} - 连接失败: {e}{self.RESET}")
        
        if not working_urls:
            print(f"{self.RED}❌ 所有镜像源都无法连接，请检查网络设置{self.RESET}")
            return False
        
        print(f"{self.GREEN}✅ 网络连接测试完成，{len(working_urls)} 个镜像源可用{self.RESET}")
        return True

    def setup_pip_mirrors(self):
        """配置pip镜像源"""
        print(f"{self.BLUE}[配置pip镜像源]{self.RESET}")
        print("-" * 45)
        
        # 镜像源配置
        primary_mirror = "https://mirrors.cloud.tencent.com/pypi/simple/"
        extra_mirrors = "https://mirrors.aliyun.com/pypi/simple/ https://pypi.tuna.tsinghua.edu.cn/simple/"
        trusted_hosts = "mirrors.cloud.tencent.com mirrors.aliyun.com pypi.tuna.tsinghua.edu.cn"
        
        # 清除旧配置
        self.run_command(f'"{self.python_exe}" -m pip config unset global.index-url', check=False)
        self.run_command(f'"{self.python_exe}" -m pip config unset global.extra-index-url', check=False)
        
        # 设置主镜像源
        print(f"主镜像源： {primary_mirror}")
        result = self.run_command(f'"{self.python_exe}" -m pip config set global.index-url "{primary_mirror}"', check=False)
        if result and result.returncode != 0:
            print(f"{self.RED}[错误] 主镜像源配置失败{self.RESET}")
            print("尝试备用方案...")
            self.run_command(f'"{self.python_exe}" -m pip install --index-url "{primary_mirror}" pip', check=False)
        
        # 设置备用镜像源
        print(f"备用镜像： {extra_mirrors}")
        self.run_command(f'"{self.python_exe}" -m pip config set global.extra-index-url "{extra_mirrors}"', check=False)
        
        # 设置信任主机
        print(f"信任域名: {trusted_hosts}")
        self.run_command(f'"{self.python_exe}" -m pip config set global.trusted-host "{trusted_hosts}"', check=False)
        
        # 验证配置
        print("-" * 45)
        print(f"{self.BLUE}[当前pip配置]{self.RESET}")
        self.run_command(f'"{self.python_exe}" -m pip config list', check=False)
        print("-" * 45)
        
        # 测试镜像源
        self.test_mirrors([primary_mirror] + extra_mirrors.split())

    def test_mirrors(self, mirrors):
        """测试镜像源连接"""
        print(f"\n{self.BLUE}[测试镜像源连接]{self.RESET}")
        for mirror in mirrors:
            print(f"测试镜像: {mirror}")
            try:
                response = requests.get(mirror, timeout=10)
                print(f"  HTTP状态: {response.status_code}")
                print(f"  响应时间: {response.elapsed.total_seconds():.2f}s")
            except Exception as e:
                print(f"  连接失败: {e}")

    def detect_cuda(self):
        """检测CUDA版本"""
        print(f"{self.CYAN}🧠 检测是否有独显...{self.RESET}")
        try:
            result = self.run_command("nvidia-smi", check=False, show_output=False)
            if result is not None and result.returncode == 0:
                print(f"{self.GREEN}✅ 检测到 NVIDIA 显卡{self.RESET}")
                
                # 解析CUDA版本
                if result.stdout:
                    for line in result.stdout.split('\n'):
                        if "CUDA Version" in line:
                            version = line.split("CUDA Version:")[1].strip().split()[0]
                            print(f"{self.GREEN}✅ 当前 CUDA 版本: {version}{self.RESET}")
                            
                            # 解析版本号
                            major, minor = version.split('.')[:2]
                            cuda_ver_num = int(major + minor)
                            
                            # 选择PyTorch版本
                            if cuda_ver_num >= 126:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu126/"
                            elif cuda_ver_num >= 124:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu124/"
                            elif cuda_ver_num >= 121:
                                return "2.5.1", "https://mirrors.aliyun.com/pytorch-wheels/121/"
                            elif cuda_ver_num >= 118:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu118/"
                            else:
                                return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"
                
                # 如果没有找到CUDA版本信息，使用默认版本
                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu118/"
            else:
                print(f"{self.YELLOW}🧊 未检测到独显，使用 CPU 模式安装 PyTorch...{self.RESET}")
                return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"
        except Exception as e:
            print(f"{self.YELLOW}🧊 检测独显时出错: {e}，使用 CPU 模式安装 PyTorch...{self.RESET}")
            return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"

    def install_packages(self, packages, extra_args=""):
        """安装Python包"""
        if packages is None:
            packages = []
        
        failed_packages = []
        
        for i, package in enumerate(packages, 1):
            print(f"\n{self.CYAN}[{i}/{len(packages)}] 安装：{package}{self.RESET}")
            cmd = f'"{self.python_exe}" -m pip install "{package}" {extra_args}'
            result = self.run_command(cmd, check=False, show_output=True)
            
            if result is None or result.returncode != 0:
                print(f"{self.YELLOW}⚠️ 安装失败，尝试安装最新版本...{self.RESET}")
                package_name = package.split('==')[0].split('>=')[0].split('<=')[0].split('~=')[0]
                retry_cmd = f'"{self.python_exe}" -m pip install "{package_name}" --upgrade {extra_args}'
                retry_result = self.run_command(retry_cmd, check=False, show_output=True)
                
                if retry_result is None or retry_result.returncode != 0:
                    print(f"{self.RED}❌ 重试安装也失败{self.RESET}")
                    failed_packages.append(package)
                else:
                    print(f"{self.GREEN}✅ 重试安装成功{self.RESET}")
            else:
                print(f"{self.GREEN}✅ 安装成功{self.RESET}")
        
        # 报告失败的包
        if failed_packages:
            print(f"\n{self.RED}❌ 以下包安装失败：{self.RESET}")
            for pkg in failed_packages:
                print(f"  - {pkg}")
            print(f"\n{self.YELLOW}💡 建议手动安装这些包或检查网络连接{self.RESET}")
        
        return len(failed_packages) == 0

    def merge_requirements(self):
        """合并requirements.txt文件"""
        requirements = []
        
        # 读取ComfyUI的requirements.txt
        comfyui_req = self.comfyui_path / "requirements.txt"
        if comfyui_req.exists():
            with open(comfyui_req, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        requirements.append(line)
        
        # 读取Manager的requirements.txt
        manager_req = self.manager_path / "requirements.txt"
        if manager_req.exists():
            with open(manager_req, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        requirements.append(line)
        
        return requirements

    def create_directories(self):
        """创建模型目录"""
        model_dirs = [
            "models", "models/checkpoints", "models/clip", "models/clip_vision",
            "models/configs", "models/controlnet", "models/diffusers",
            "models/embeddings", "models/gligen", "models/hypernetworks",
            "models/loras", "models/upscale_models", "models/vae",
            "models/vae_approx", "models/custom_nodes", "models/ultralytics"
        ]
        
        for dir_name in model_dirs:
            dir_path = self.config_path / dir_name
            dir_path.mkdir(parents=True, exist_ok=True)
        
        print(f"{self.GREEN}✅ 创建模型目录结构{self.RESET}")

    def create_extra_model_paths(self):
        """创建extra_model_paths.yaml配置文件"""
        config_content = f"""# ComfyUI 额外模型路径配置
# 自动生成于 {time.strftime('%Y-%m-%d %H:%M:%S')}

# 主模型目录
models: {self.config_path}/models

# 检查点模型
checkpoints: {self.config_path}/models/checkpoints
configs: {self.config_path}/models/configs

# 文本编码器
clip: {self.config_path}/models/clip
clip_vision: {self.config_path}/models/clip_vision

# 控制网络
controlnet: {self.config_path}/models/controlnet

# 扩散模型
diffusers: {self.config_path}/models/diffusers

# 嵌入模型
embeddings: {self.config_path}/models/embeddings

# GLIGEN模型
gligen: {self.config_path}/models/gligen

# 超网络
hypernetworks: {self.config_path}/models/hypernetworks

# LoRA模型
loras: {self.config_path}/models/loras

# 上采样模型
upscale_models: {self.config_path}/models/upscale_models

# VAE模型
vae: {self.config_path}/models/vae
vae_approx: {self.config_path}/models/vae_approx

# 自定义节点
custom_nodes: {self.config_path}/models/custom_nodes

# Ultralytics模型
ultralytics: {self.config_path}/models/ultralytics
"""
        
        config_file = self.comfyui_path / "extra_model_paths.yaml"
        with open(config_file, 'w', encoding='utf-8') as f:
            f.write(config_content)
        print(f"{self.GREEN}✅ 生成 extra_model_paths.yaml 配置文件{self.RESET}")

    def create_manager_configs(self):
        """创建ComfyUI-Manager配置文件"""
        # 创建目录
        manager_user_dir = self.comfyui_path / "user" / "default" / "ComfyUI-Manager"
        manager_user_dir.mkdir(parents=True, exist_ok=True)
        
        manager_config_dir = self.config_path / "custom_nodes" / "ComfyUI-Manager"
        manager_config_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成channels.list
        channels_content = """default::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main
recent::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main/node_db/new
legacy::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main/node_db/legacy
forked::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main/node_db/forked
dev::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main/node_db/dev
tutorial::https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main/node_db/tutorial
"""
        
        channels_file = manager_user_dir / "channels.list"
        with open(channels_file, 'w', encoding='utf-8') as f:
            f.write(channels_content)
        
        # 复制到config目录
        shutil.copy2(channels_file, manager_config_dir / "channels.list")
        
        # 生成config.ini
        config_ini_content = """[default]
preview_method = none
badge_mode = none
git_exe = 
channel_url = https://gitee.com/chunfengruge/ComfyUI-Manager/raw/main
share_option = all
bypass_ssl = False
file_logging = True
default_ui = none
component_policy = workflow
double_click_policy = copy-all
windows_selector_event_loop_policy = False
model_download_by_agent = False
downgrade_blacklist = 
security_level = normal
skip_migration_check = False
"""
        
        config_ini_file = manager_user_dir / "config.ini"
        with open(config_ini_file, 'w', encoding='utf-8') as f:
            f.write(config_ini_content)
        
        print(f"{self.GREEN}✅ 生成ComfyUI-Manager配置文件{self.RESET}")

    def create_launcher_script(self, extra_args=""):
        """创建启动脚本"""
        launcher_content = f"""@echo off
chcp 65001 >nul
title ComfyUI启动器
setlocal enabledelayedexpansion

set "WORKSPACE={self.workspace_path}"
set "COMFYUI_PATH={self.comfyui_path}"

set "PYTHON_DIR=!WORKSPACE!\\python"
set "PATH=!PYTHON_DIR!;!PYTHON_DIR!\\py_scripts;!PATH!"
set "PYTHONPATH=!PYTHON_DIR!\\Lib\\site-packages"
set "PYTHON=!WORKSPACE!\\python\\python.exe"

:: 配置pip镜像源
set "PRIMARY_MIRROR=https://mirrors.cloud.tencent.com/pypi/simple/"
set "EXTRA_MIRRORS=https://mirrors.aliyun.com/pypi/simple/ https://pypi.tuna.tsinghua.edu.cn/simple/"
set "TRUSTED_HOSTS=mirrors.cloud.tencent.com mirrors.aliyun.com pypi.tuna.tsinghua.edu.cn"

!PYTHON! -m pip config set global.index-url "!PRIMARY_MIRROR!" >nul
!PYTHON! -m pip config set global.extra-index-url "!EXTRA_MIRRORS!" >nul
!PYTHON! -m pip config set global.trusted-host "!TRUSTED_HOSTS!" >nul

:: 设置HuggingFace镜像
set HF_ENDPOINT=https://hf-mirror.com

:: 配置Git
git config --global http.sslVerify false
git config --global url."https://gitclone.com/github.com/".insteadOf https://github.com/

set "PORT=8188"

:: 检查端口占用
for /f "tokens=5 delims= " %%P in ('netstat -ano ^| findstr ":!PORT!"') do (
   taskkill /f /pid %%P >nul
   echo [INFO] 已终止占用端口!PORT!的进程(PID: %%P)
)

if exist "!COMFYUI_PATH!\\main.py" (
   !PYTHON! "!COMFYUI_PATH!\\main.py" {extra_args} --port !PORT! || (
       echo [ERROR] 启动失败
       pause
       exit /b 1
   )
) else (
   echo [ERROR] main.py未找到: !COMFYUI_PATH!\\main.py
   pause
   exit /b 1
)

pause
"""
        
        launcher_file = self.workspace_path / "ComfyUI.bat"
        with open(launcher_file, 'w', encoding='utf-8') as f:
            f.write(launcher_content)
        
        print(f"{self.GREEN}✅ 生成启动脚本 ComfyUI.bat{self.RESET}")

    def install_comfyui(self) -> bool:
        """执行ComfyUI安装流程"""
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        print(f"{self.PURPLE}ComfyUI 安装脚本 (插件版本){self.RESET}")
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        
        # 测试网络连接
        if not self.test_network_connection():
            return False
        
        # 克隆ComfyUI
        print(f"克隆 ComfyUI 主体到 {self.comfyui_path}")
        if not self.clone_repository(
            "https://gitcode.com/chunfengruge/ComfyUI.git",
            self.comfyui_path,
            "ComfyUI"
        ):
            return False
        
        # 克隆ComfyUI-Manager
        print(f"{self.CYAN}🧩 安装 ComfyUI-Manager 插件...{self.RESET}")
        if not self.clone_repository(
            "https://gitcode.com/chunfengruge/ComfyUI-Manager.git",
            self.manager_path,
            "ComfyUI-Manager"
        ):
            return False
        
        # 配置pip镜像源
        self.setup_pip_mirrors()
        
        # 升级pip
        print(f"{self.CYAN}📦 升级pip...{self.RESET}")
        self.run_command(f'"{self.python_exe}" -m pip install --upgrade pip', check=False, show_output=True)
        
        # 安装核心依赖
        print(f"{self.CYAN}📦 安装核心依赖...{self.RESET}")
        self.run_command(f'"{self.python_exe}" -m pip install diffusers transformers', check=False, show_output=True)
        
        # 检测CUDA并安装PyTorch
        torch_ver, torch_url = self.detect_cuda()
        print(f"{self.GREEN}✅ 自动选择的 PyTorch 安装版本: {torch_ver}{self.RESET}")
        print(f"{self.CYAN}🌐 安装源: {torch_url}{self.RESET}")
        
        if torch_ver == "cpu":
            print(f"{self.CYAN}📦 安装CPU版本PyTorch...{self.RESET}")
            self.run_command(f'"{self.python_exe}" -m pip install torch==2.6.0 torchvision torchaudio --retries 5', check=False, show_output=True)
        else:
            print(f"{self.CYAN}📦 安装GPU版本PyTorch...{self.RESET}")
            self.run_command(f'"{self.python_exe}" -m pip install torch=={torch_ver} torchvision torchaudio -f {torch_url} --retries 5', check=False, show_output=True)
        
        # 合并并安装requirements.txt
        print(f"{self.CYAN}🔄 合并安装依赖...{self.RESET}")
        requirements = self.merge_requirements()
        if requirements is None:
            print(f"{self.YELLOW}⚠️ 未找到requirements.txt文件，跳过依赖安装{self.RESET}")
            requirements = []
        else:
            print(f"{self.CYAN}📋 找到 {len(requirements)} 个依赖包{self.RESET}")
        
        self.install_packages(requirements)
        
        # 创建目录结构
        print(f"{self.CYAN}📁 创建目录结构...{self.RESET}")
        self.create_directories()
        
        # 生成配置文件
        print(f"{self.CYAN}⚙ 生成配置文件...{self.RESET}")
        self.create_extra_model_paths()
        self.create_manager_configs()
        
        # 创建启动脚本
        print(f"{self.CYAN}💡 正在生成启动脚本...{self.RESET}")
        extra_args = "" if torch_ver != "cpu" else "--cpu"
        self.create_launcher_script(extra_args)
        
        # 安装完成
        print(f"\n{self.BLUE}{'=' * 60}{self.RESET}")
        print(f"{self.GREEN}✅ 安装完成！双击 ComfyUI.bat 启动 ComfyUI 吧～{self.RESET}")
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        
        input("按回车键退出...")
        return True

def main():
    """主函数"""
    installer = ComfyUIInstaller()
    return installer.start()

if __name__ == "__main__":
    sys.exit(0 if main() else 1) 